import os

import requests
from dotenv import load_dotenv
import numpy as np
import pandas as pd
from langchain.chat_models import init_chat_model
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate, \
    PromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableLambda
from langchain_experimental.tools.python.tool import PythonAstREPLTool
from langchain_core.output_parsers import JsonOutputKeyToolsParser
from langchain_core.tools import tool as chain_tool
from langchain.agents import create_tool_calling_agent,tool as agent_tool
from langchain.agents import AgentExecutor
import json


# 测试pandas读取csv文件
def test_read_csv():
    dataset = pd.read_csv("./global_cities_data.csv")
    pd.set_option('max_colwidth', 200)
    print(dataset.head(5))
    dataset.info()

# 测试langchain内置python解释器工具，可用于运行python代码，需安装依赖pip3 install langchain_experimental
def test_langchain_built_in_python_interpreter():
    df = pd.read_csv("./global_cities_data.csv")
    # 这是对df得定义，即定义局部变量
    tool = PythonAstREPLTool(locals={"df": df})
    # invoke执行得是一个python代码，返回值为代码执行结果，此处为获取该列均值
    invoke_result = tool.invoke("df['GDP_Billion_USD'].mean()")
    print("!!! Python Interpreter Tool invoke result:", invoke_result)

# 测试langchain在链中使用内置python解释器工具
def test_chain_with_built_in_tool():
    def debug_print(text):
        print(f"!!产生的python代码是: {text}")
        return text
    debug_node = RunnableLambda(debug_print)
    system_prompt = "你可以访问一个名为 `df` 的 pandas 数据框，请根据用户提出的问题，编写 Python 代码来回答。只返回代码，不返回其他内容。只允许使用 pandas 和内置库。"
    human_prompt = """
            这是用户的问题： {user_query}
            """
    system_prompt_template = SystemMessagePromptTemplate.from_template(system_prompt)
    human_prompt_template = HumanMessagePromptTemplate.from_template(human_prompt)
    prompt_template = ChatPromptTemplate.from_messages([system_prompt_template, human_prompt_template])


    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    df = pd.read_csv("./global_cities_data.csv")
    # 这是对df的定义，即定义局部变量
    tool = PythonAstREPLTool(locals={"df": df})
    # 将tool绑定给大语言模型
    model = model.bind_tools([tool])
    # 将关键部分提取出
    parser = JsonOutputKeyToolsParser(key_name=tool.name,first_tool_only=True)
    chain = prompt_template|  model  | parser |debug_node |tool
    result = chain.invoke({"user_query":"我有一张表，名为'df'，请帮我计算GDP_Billion_USD列的均值。"} )
    print(f"!!! Langchain with tool, result:{result}")

# 使用tool装饰器自定义langchain的外部工具，必须为工具指定description，即""" ... """，且必须包含参数说明
@chain_tool()
@agent_tool()
def get_weather(location):
    """
    获取指定位置的当前天气信息
    :param location: 位置是中国城市名，例如：上海
    :return: OpenWeather API查询即时天气的结果， 返回结果对象类型为解析之后的JSON格式对象，并用字符串形式进行表示，其中包含了全部重要的天气信息
    """

    # Step 1.构建请求
    url = "https://getweather.api.bdymkt.com/lundear/weather1d"

    # Step 2.设置查询参数
    params = {
        "areaCn": location,
    }

    # 添加header
    headers = {
        "X-Bce-Signature": f"AppCode/{os.getenv('BAIDU_YUN_API_APP_CODE')}"
    }


    # Step 3.发送GET请求
    response = requests.get(url, params=params,headers=headers)

    # Step 4.解析响应
    data = response.json()
    return json.dumps(data)

# langchain接入自定义工具
def test_chain_with_custom_tool():
    def debug_print(text):
        print(f"!!中间过程: {text}")
        return text

    # 初始化大语言模型
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    # 将自定义工具绑定给大语言模型
    model = model.bind_tools([get_weather])
    # 定义json解析工具，抓取目标工具的输出
    paser = JsonOutputKeyToolsParser(key_name=get_weather.name,first_tool_only=True)
    # 定义调用获取天气api的chain, 此链用于llm识别用户输入，理解要调用的工具，并生成调用该工具的参数，然后调用工具获得结果
    get_weather_chain = model | paser | debug_print | get_weather


    # 定义第二个链，用于将api json 结果转换为自然语言
    output_system_prompt =PromptTemplate.from_template("""
    你将收到一段 JSON 格式的天气数据，请用简洁自然的方式将其转述给用户。
    以下是天气 JSON 数据：
    
    {weather_json}
    
    请将其转换为中文天气描述，例如：
    “北京当前天气晴，气温为 23°C，湿度 58%，风速 2.1 米/秒。”
    只返回一句话描述，不要其他说明或解释。
    """)

    output_chain = output_system_prompt | model | StrOutputParser()

    # 将两个链组合成一个完整的链
    full_chain = get_weather_chain | output_chain
    response = full_chain.invoke("成都天气怎么样？")
    print(f"!!! Langchain with custom tool, result:{response}")


@agent_tool()
def write_file(content):
    """
    将指定内容写入本地文件。
    :param content: 必要参数，字符串类型，用于表示需要写入文档的具体内容。
    :return：是否成功写入
    """
    with open("output.txt", "w", encoding="utf-8") as f:
        f.write(content)

    return "已成功写入本地文件。"

# LangChain Agent使用function calling方式接入自定义工具
def test_agent_with_custom_tool_by_function_calling(input):
    prompt = ChatPromptTemplate.from_messages([
        SystemMessagePromptTemplate.from_template("你是天气助手，请根据用户的问题，给出相应的天气信息"),
        HumanMessagePromptTemplate.from_template("{input}"),
        MessagesPlaceholder(variable_name="agent_scratchpad"),
    ])

    tools = [get_weather]
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    model.bind_tools(tools)
    # 创建langchain agent
    agent = create_tool_calling_agent(model, tools, prompt)
    # 使用agent executor来执行agent, verbose=True会打印中间过程
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
    response = agent_executor.invoke({"input": input})
    print(f"!!! Langchain agent with custom tool by function calling, response:{response}, output:{response['output']}")


# LangChain Agent使用function calling方式接入多个自定义工具实现串行调用
def test_agent_with_multi_custom_tool_by_function_calling(input):
    prompt = ChatPromptTemplate.from_messages([
        SystemMessagePromptTemplate.from_template("你是天气助手，请根据用户的问题，给出相应的天气信息，如果用户需要保存此结果，请写入到文件"),
        HumanMessagePromptTemplate.from_template("{input}"),
        MessagesPlaceholder(variable_name="agent_scratchpad"),
    ])

    tools = [get_weather,write_file]
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    model.bind_tools(tools)
    # 创建langchain agent
    agent = create_tool_calling_agent(model, tools, prompt)
    # 使用agent executor来执行agent, verbose=True会打印中间过程
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
    response = agent_executor.invoke({"input": input})
    print(f"!!! Langchain agent with custom tool by function calling, response:{response}, output:{response['output']}")



if __name__ == '__main__':
    load_dotenv(override=True)
    DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
    print(f"DEEPSEEK_API_KEY: {DEEPSEEK_API_KEY}")

    # test_read_csv()
    # test_langchain_built_in_python_interpreter()
    # test_chain_with_built_in_tool()

    # print(f"get_weather tool name:{get_weather.name}")
    # print(f"get_weather tool desc:{get_weather.description}")
    # print(f"get_weather tool args:{get_weather.args}")

    # test_chain_with_custom_tool()
    # test_agent_with_custom_tool_by_function_calling(input="成都的天气怎么样")

    # agent支持处理复杂的问题，对于该提问他会并行调用工具两次，分别查询成都和上海的天气，然后根据用户的问题进行回答
    # test_agent_with_custom_tool_by_function_calling(input="成都和上海的天气怎么样，哪个地方更热")
    test_agent_with_multi_custom_tool_by_function_calling(input="成都和上海的天气怎么样，哪个地方更热，把详细的比较结果保存到文件呢")





